"""
扫描操作符实现
"""

from typing import Optional, List, Dict, Any, Iterator
import time
import threading

from ..executor_base import IteratorExecutor
from ..types import (
    Record, RecordBatch, ExecutionContext, OperatorType,
    TableInfo, IndexInfo, DataValue
)
from ..storage_adapter import StorageEngineAdapter

class TableScanOperator(IteratorExecutor):
    """表扫描操作符"""
    
    def __init__(self, executor_id: str, context: ExecutionContext, 
                 table_name: str, table_info: TableInfo = None,
                 storage_engine: StorageEngineAdapter = None,
                 table_alias: str = None):
        super().__init__(executor_id, context)
        self.table_name = table_name
        self.table_alias = table_alias or table_name  # 设置表别名
        self.table_info = table_info
        self.storage_engine = storage_engine
        self.scan_position = 0
        self.total_rows = 0
        self.predicate_filters: List[callable] = []
        self.projection_columns: Optional[List[str]] = None
        self.predicate_string: Optional[str] = None
        
        # 如果没有提供存储引擎，使用模拟数据
        self.use_mock_data = storage_engine is None
        if self.use_mock_data:
            self.table_data: List[Record] = []
            self._initialize_sample_data()
        else:
            # 获取表信息
            if not self.table_info:
                self.table_info = self.storage_engine.get_table_info(table_name)
    
    def _initialize_sample_data(self):
        """初始化示例数据"""
        # 生成一些示例数据用于测试
        if self.table_name == "users":
            self.table_data = [
                {"id": 1, "name": "张三", "age": 25, "city": "北京"},
                {"id": 2, "name": "李四", "age": 30, "city": "上海"},
                {"id": 3, "name": "王五", "age": 28, "city": "广州"},
                {"id": 4, "name": "赵六", "age": 35, "city": "深圳"},
                {"id": 5, "name": "钱七", "age": 22, "city": "杭州"},
            ]
        elif self.table_name == "orders":
            self.table_data = [
                {"order_id": 1001, "user_id": 1, "amount": 199.99, "status": "completed"},
                {"order_id": 1002, "user_id": 2, "amount": 299.50, "status": "pending"},
                {"order_id": 1003, "user_id": 1, "amount": 89.00, "status": "completed"},
                {"order_id": 1004, "user_id": 3, "amount": 450.75, "status": "shipped"},
                {"order_id": 1005, "user_id": 4, "amount": 120.30, "status": "completed"},
            ]
        elif self.table_name == "student":
            # 添加学生表数据用于测试
            self.table_data = [
                {"id": 1, "name": "张三", "age": 18, "major": "计算机科学"},
                {"id": 2, "name": "李四", "age": 22, "major": "软件工程"},
                {"id": 3, "name": "王五", "age": 21, "major": "数据科学"},
                {"id": 4, "name": "赵六", "age": 19, "major": "人工智能"},
                {"id": 5, "name": "钱七", "age": 23, "major": "计算机科学"},
                {"id": 6, "name": "孙八", "age": 20, "major": "软件工程"},
                {"id": 7, "name": "周九", "age": 24, "major": "网络安全"},
                {"id": 8, "name": "吴十", "age": 22, "major": "数据科学"},
            ]
        else:
            # 默认生成一些通用数据
            self.table_data = [
                {"id": i, "value": f"data_{i}", "timestamp": time.time() + i}
                for i in range(1, 101)
            ]
        
        self.total_rows = len(self.table_data)
    
    def add_predicate_filter(self, predicate: callable):
        """添加谓词过滤器"""
        self.predicate_filters.append(predicate)
    
    def set_predicate_string(self, predicate: str):
        """设置谓词字符串（用于存储引擎查询优化）"""
        self.predicate_string = predicate
    
    def set_projection(self, columns: List[str]):
        """设置投影列"""
        self.projection_columns = columns
    
    def open(self):
        """打开扫描操作符"""
        with self.lock:
            if self.is_open:
                return
            
            self.is_open = True
            self.scan_position = 0
            self.current_iterator = self._create_scan_iterator()
            
            # 更新统计信息
            self.update_stats(io_ops=1)
    
    def _create_scan_iterator(self) -> Iterator[Record]:
        """创建扫描迭代器"""
        if self.use_mock_data:
            # 使用模拟数据
            for record in self.table_data:
                # 检查超时
                self.check_timeout()
                
                # 应用谓词过滤
                if self._apply_predicates(record):
                    # 应用投影
                    projected_record = self._apply_projection(record)
                    
                    # 更新统计
                    self.update_stats(rows_processed=1)
                    
                    yield projected_record
        else:
            # 使用真实存储引擎
            try:
                for record in self.storage_engine.scan_table(
                    self.table_name, 
                    columns=self.projection_columns,
                    predicate=self.predicate_string
                ):
                    # 检查超时
                    self.check_timeout()
                    
                    # 应用内存中的谓词过滤（用于存储引擎无法处理的复杂条件）
                    if self._apply_predicates(record):
                        # 更新统计
                        self.update_stats(rows_processed=1)
                        
                        yield record
            except Exception as e:
                # 如果存储引擎出错，记录错误并停止扫描
                print(f"Storage engine error during table scan: {e}")
                return
    
    def _apply_predicates(self, record: Record) -> bool:
        """应用谓词过滤"""
        for predicate in self.predicate_filters:
            try:
                if not predicate(record):
                    return False
            except Exception:
                # 谓词执行出错，过滤掉该记录
                return False
        return True
    
    def _apply_projection(self, record: Record) -> Record:
        """应用投影"""
        if self.projection_columns is None:
            return record
        
        projected = {}
        for col in self.projection_columns:
            if col in record:
                projected[col] = record[col]
            else:
                projected[col] = None
        
        return projected
    
    def close(self):
        """关闭扫描操作符"""
        with self.lock:
            if not self.is_open:
                return
            
            self.is_open = False
            self.current_iterator = None
            
            # 更新最终统计
            self.stats.execution_time_ms = (time.time() - self.context.start_time) * 1000
    
    def get_estimated_cost(self) -> float:
        """获取估算成本"""
        # 简单的成本模型：基于表大小
        base_cost = self.total_rows * 0.1  # 每行0.1个成本单位
        
        # 谓词过滤可以降低成本
        if self.predicate_filters:
            base_cost *= 0.5
        
        return base_cost
    
    def get_estimated_rows(self) -> int:
        """获取估算行数"""
        if self.use_mock_data:
            estimated = self.total_rows
            
            # 简单的选择性估算
            if self.predicate_filters:
                estimated = int(estimated * 0.3)  # 假设过滤后剩余30%
            
            return max(1, estimated)
        else:
            # 从存储引擎获取统计信息
            try:
                # 简单估算：假设表有1000行，过滤后剩余30%
                estimated = 1000
                if self.predicate_filters or self.predicate_string:
                    estimated = int(estimated * 0.3)
                return max(1, estimated)
            except:
                return 1000

class IndexScanOperator(IteratorExecutor):
    """索引扫描操作符"""
    
    def __init__(self, executor_id: str, context: ExecutionContext,
                 table_name: str, index_info: IndexInfo,
                 scan_range: Optional[Dict[str, Any]] = None):
        super().__init__(executor_id, context)
        self.table_name = table_name
        self.index_info = index_info
        self.scan_range = scan_range or {}
        self.current_position = 0
        
        # 模拟索引数据
        self.index_entries: List[Dict[str, Any]] = []
        self._initialize_index_data()
    
    def _initialize_index_data(self):
        """初始化索引数据"""
        # 模拟索引项，包含键值和行ID
        if self.table_name == "users" and "id" in self.index_info.columns:
            self.index_entries = [
                {"key": {"id": 1}, "row_id": 1, "row_data": {"id": 1, "name": "张三", "age": 25}},
                {"key": {"id": 2}, "row_id": 2, "row_data": {"id": 2, "name": "李四", "age": 30}},
                {"key": {"id": 3}, "row_id": 3, "row_data": {"id": 3, "name": "王五", "age": 28}},
                {"key": {"id": 4}, "row_id": 4, "row_data": {"id": 4, "name": "赵六", "age": 35}},
                {"key": {"id": 5}, "row_id": 5, "row_data": {"id": 5, "name": "钱七", "age": 22}},
            ]
        else:
            # 默认索引数据
            self.index_entries = [
                {"key": {"id": i}, "row_id": i, "row_data": {"id": i, "value": f"data_{i}"}}
                for i in range(1, 11)
            ]
    
    def set_scan_range(self, start_key: Dict[str, Any] = None, 
                      end_key: Dict[str, Any] = None,
                      inclusive_start: bool = True,
                      inclusive_end: bool = True):
        """设置扫描范围"""
        self.scan_range = {
            'start_key': start_key,
            'end_key': end_key,
            'inclusive_start': inclusive_start,
            'inclusive_end': inclusive_end
        }
    
    def open(self):
        """打开索引扫描操作符"""
        with self.lock:
            if self.is_open:
                return
            
            self.is_open = True
            self.current_position = 0
            
            # 根据扫描范围过滤索引项
            filtered_entries = self._filter_by_range()
            self.current_iterator = self._create_index_iterator(filtered_entries)
            
            # 更新统计信息
            self.update_stats(io_ops=1)
    
    def _filter_by_range(self) -> List[Dict[str, Any]]:
        """根据范围过滤索引项"""
        if not self.scan_range:
            return self.index_entries
        
        filtered = []
        start_key = self.scan_range.get('start_key')
        end_key = self.scan_range.get('end_key')
        inclusive_start = self.scan_range.get('inclusive_start', True)
        inclusive_end = self.scan_range.get('inclusive_end', True)
        
        for entry in self.index_entries:
            key = entry['key']
            
            # 检查起始范围
            if start_key is not None:
                if not self._key_compare(key, start_key, inclusive_start, True):
                    continue
            
            # 检查结束范围
            if end_key is not None:
                if not self._key_compare(key, end_key, inclusive_end, False):
                    continue
            
            filtered.append(entry)
        
        return filtered
    
    def _key_compare(self, key: Dict[str, Any], boundary_key: Dict[str, Any],
                    inclusive: bool, is_start: bool) -> bool:
        """键值比较"""
        # 简化的键值比较逻辑
        for col in self.index_info.columns:
            if col in key and col in boundary_key:
                key_val = key[col]
                boundary_val = boundary_key[col]
                
                if key_val < boundary_val:
                    return not is_start
                elif key_val > boundary_val:
                    return is_start
                # key_val == boundary_val，检查包含性
        
        return inclusive
    
    def _create_index_iterator(self, entries: List[Dict[str, Any]]) -> Iterator[Record]:
        """创建索引扫描迭代器"""
        for entry in entries:
            # 检查超时
            self.check_timeout()
            
            # 返回行数据
            record = entry.get('row_data', {})
            
            # 更新统计
            self.update_stats(rows_processed=1, io_ops=1)
            
            yield record
    
    def close(self):
        """关闭索引扫描操作符"""
        with self.lock:
            if not self.is_open:
                return
            
            self.is_open = False
            self.current_iterator = None
            
            # 更新最终统计
            self.stats.execution_time_ms = (time.time() - self.context.start_time) * 1000
    
    def get_estimated_cost(self) -> float:
        """获取估算成本"""
        # 索引扫描通常比表扫描更高效
        base_cost = len(self.index_entries) * 0.05  # 每项0.05个成本单位
        
        # 范围扫描可能更高效
        if self.scan_range:
            base_cost *= 0.3
        
        return base_cost
    
    def get_estimated_rows(self) -> int:
        """获取估算行数"""
        if not self.scan_range:
            return len(self.index_entries)
        
        # 简单估算：范围扫描返回20%的数据
        return max(1, int(len(self.index_entries) * 0.2))

class SeekableTableScanOperator(TableScanOperator):
    """可定位的表扫描操作符"""
    
    def __init__(self, executor_id: str, context: ExecutionContext,
                 table_name: str, table_info: TableInfo = None):
        super().__init__(executor_id, context, table_name, table_info)
        self.bookmarks: Dict[str, int] = {}  # 书签位置
    
    def seek_to_position(self, position: int):
        """定位到指定位置"""
        if 0 <= position < len(self.table_data):
            self.scan_position = position
    
    def create_bookmark(self, bookmark_name: str):
        """创建书签"""
        self.bookmarks[bookmark_name] = self.scan_position
    
    def seek_to_bookmark(self, bookmark_name: str) -> bool:
        """定位到书签位置"""
        if bookmark_name in self.bookmarks:
            self.scan_position = self.bookmarks[bookmark_name]
            return True
        return False
    
    def get_current_position(self) -> int:
        """获取当前位置"""
        return self.scan_position

